# encoding: utf-8

"""
模块描述
使用pandas构造tensorflow的输入样本

Authors: tongzhenguo
Date:    2023/01/19
"""
import numpy as np
from pandas import DataFrame
from sklearn.preprocessing import MinMaxScaler


def gen_windows(X_in, y_in, window_size):
    X_out = []
    y_out = []
    for i in range(window_size, X_in.shape[0]):
        X_out.append(X_in[i - window_size:i, ])
        y_out.append(y_in[i - 1])
    return np.array(X_out), np.array(y_out)


class Sample(object):
    def __init__(self):
        # 训练集分割占比
        self.train_test_split = 0.8
        # 样本的时间跨度T,T-1个时间作为
        self.window_size = 5
        # 输入特征的时间窗口
        self.input_window_size = 1

    def generate(self, df: DataFrame):
        """滑动窗口构造时序样本
        注意：默认输入的df是时间升序的
        """
        # 特征生成
        df['previous_close'] = df['close'].shift(self.input_window_size)
        # 删除掉previous_close为nan的数据
        df = df.dropna()
        # 数据预处理
        columns = ['previous_close', 'close']
        self.scaler = MinMaxScaler(feature_range=(0, 1)).fit(df[columns].values)
        # min,max
        print(self.scaler.data_min_)
        print(self.scaler.data_max_)
        df[columns] = self.scaler.transform(df[columns].values)
        # 定义特征、target
        X = df[['previous_close']].values
        y = df['close'].values
        # 训练集测试集分割
        split = int(X.shape[0] * self.train_test_split)
        train_x, train_y = X[:split], y[:split]
        test_x, test_y = X[split:], y[split:]
        train_x, train_y = gen_windows(train_x, train_y, self.window_size)
        test_x, test_y = gen_windows(test_x, test_y, self.window_size)
        return train_x, train_y, test_x, test_y
